#include "refine/refine3.h"
#include "imglib/imageclass.h"
#include "imglib/bitmap.h"
#include "imglib/bmpmorph.h"
#include <iostream>
#include <limits>
#include <list>
#include <math.h>
using std::clog;
using std::endl;

namespace refine {

namespace refine3 {

params::params()
{
	t = 5.0f;
	max_passes = 150;
	consider = info_all;
	change = turn_on_and_off;
	prefer = prefer_none;
	silent = false;

	modifiable = 0;
}

bool operator == (const params & lhs, const params & rhs)
{
	if (lhs.t != rhs.t ||
		lhs.prefer != rhs.prefer ||
		lhs.max_passes != rhs.max_passes ||
		lhs.consider != rhs.consider ||
		lhs.change != rhs.change ||
		//lhs.silent != rhs.silent ||
		lhs.modifiable != rhs.modifiable)
		return false;

	return true;
}

typedef signed char schar;
typedef unsigned short ushort;
typedef unsigned int uint;

typedef point<schar> spoint; ///< relative coordinates of a neighbour pixel
typedef rect<schar> srect; 
typedef std::pair<ushort, spoint> cdist_spoint_pair; ///< memory saving is significant here
typedef std::vector<cdist_spoint_pair> neighbours;
typedef img::BGR<int> icolor;

/// The object stores coordinates of neighbours for some points,
/// which are similar to given colors, in sorted order.
/// The sorting is very time-consuming, that is why caching is required.
template <class R>
class sorted_neighbours {
public:
	sorted_neighbours(const R & im_, int whs, const params & p_);

	const R & image() const
		{ return im; }
	const refine3::params & params() const
		{ return p; }
	uint size() const ///< the number of elements in the cache (each element correspond to a pixel)
		{ return pool.size(); }
	uint scans() const ///< the total number of neighbours scans made during lifetime of the object
		{ return total_scans; }
	uint sorts() const ///< the total number of partial sorts made during lifetime of the object
		{ return total_sorts; }

	/// \param sorted_count [in/out] is the number of sorted elements in the beginning of neighbours array
	const neighbours & get(size_point p, const icolor & centre, uint & sorted_count);
	/// frees unused information during last pass, and marks remaining information as unused
	void pass_done();

private:
	typedef point<ushort> ushort_point; 
	struct cell {
		cell() { }
		cell(ushort_point r) : root(r), last_used(0) { }

		ushort_point root;
		neighbours nei;
		size_t sorted_count;
		size_t last_used;
	};
	typedef std::list<cell> precomputed_list;

private:
	cell & get_cell(size_point p);
	template <class D>
	void scan_neighbours(neighbours & n, size_point p, const icolor & centre, const D & distance);

private:
	const R & im;
	const int whs;
	const refine3::params & p;
	const uint max_sorted;
	img::simple_image<cell*> map; ///< only requested cells of the image contains valid data
	precomputed_list pool;
	uint total_scans, total_sorts;
};

template <class R>
sorted_neighbours<R>::sorted_neighbours(const R & im_, int whs_, const refine3::params & p_) :
	im(im_), whs(whs_), p(p_),
	max_sorted(std::max(7, whs_ * (whs_ + 1))), //one third of pixels in circle window
	map(im_.dim())
{
	img::fill(map, (cell*)0);
	total_scans = total_sorts = 0;
}

template <class R>
typename sorted_neighbours<R>::cell & sorted_neighbours<R>::get_cell(size_point p)
{
	cell * & c = map[p];
	if (!c) {
		pool.push_back(cell(p));
		c = &pool.back();
	}
	else {
		c->last_used = 0;
		assert(c->root == p);
	}
	return *c;
}

template <class R>
void sorted_neighbours<R>::pass_done()
{
	precomputed_list::iterator i = pool.begin();
	while (i != pool.end()) {
		if (i->last_used < 2) {
			++i->last_used;
			++i;
		}
		else {
			map[size_point(i->root)] = 0;
			i = pool.erase(i);
		}
	}
}

struct color_distance {
	unsigned operator () (const icolor & delta) const
		{ return delta.sqr<unsigned>(); }
};

struct lightness_distance {
	unsigned operator () (const icolor & delta) const
		{ return abs(delta.r + delta.g + delta.b); }
};

struct hue_distance {
	unsigned operator () (const icolor & delta) const
		{ return (square(delta.r - delta.g) + square(delta.g - delta.b) + square(delta.b - delta.r))/2; }
};

template <class R>
template <class D>
void sorted_neighbours<R>::scan_neighbours(neighbours & n, size_point p, const icolor & centre, const D & distance)
{
	++total_scans;
	assert (whs <= std::numeric_limits<schar>::max());
	int_rect wnd;
	wnd.x1 = p.x >= whs ? -whs : -(int)p.x;
    wnd.y1 = p.y >= whs ? -whs : -(int)p.y;
	wnd.x2 = p.x + whs + 1 <= im.width() ? whs : im.width() - p.x;
	wnd.y2 = p.y + whs + 1 <= im.height() ? whs : im.height() - p.y;
	size_t area = (wnd.x2 - wnd.x1) * (wnd.y2 - wnd.y1);

	n.reserve(area - 1);

	spoint d;
	for (d.y = wnd.y1; d.y < wnd.y2; ++d.y) {
		for (d.x = wnd.x1; d.x < wnd.x2; ++d.x) 
		{
			icolor ref_col = im(p.x + d.x, p.y + d.y);
			if (d.x == 0 && d.y == 0 && centre == ref_col) //centre is equal to ref_col if they are taken from the same image, but may also be equal accidentally
				continue; //it's not right way for reference image
			if (whs > 2 && (d.x*d.x + d.y*d.y) > whs * whs)
				continue; //use circle window

			unsigned sim = distance(centre - ref_col);
			if (sim > std::numeric_limits<cdist_spoint_pair::first_type>::max())
				continue;

			n.push_back(cdist_spoint_pair(sim, d));
		}
	}

	assert (n.size() <= area - 1);
}

template <class R>
const neighbours & sorted_neighbours<R>::get(size_point p, const icolor & centre, uint & sorted_count)
{
	cell & c = get_cell(p);
	neighbours & n = c.nei;
	if (n.empty()) { 
		switch (this->p.consider) {
		case info_all:
			scan_neighbours(n, p, centre, color_distance());
			break;
		case info_lightness:
			scan_neighbours(n, p, centre, lightness_distance());
			break;
		case info_color:
			scan_neighbours(n, p, centre, hue_distance());
			break;
		default:
			assert(false);
		}
		c.sorted_count = 0;
	}
	else {
		if (c.sorted_count >= sorted_count || c.sorted_count == n.size()) {
			sorted_count = c.sorted_count;
			return n;
		}
	}

	uint max_sorted = std::min(this->max_sorted, n.size());

	if (sorted_count > max_sorted)
		sorted_count = max_sorted;
	std::partial_sort(n.begin(), n.begin() + sorted_count, n.end());
	++total_sorts;
	c.sorted_count = sorted_count;
	if (sorted_count == max_sorted && max_sorted < n.size()) {
		//remove tail of unsorted elements
		n.swap(neighbours(n.begin(), n.begin() + max_sorted));
	}

	return n;
}

////////////////////
// Core functions //
////////////////////

enum classification {
	class_0,
	class_1,
	class_denial
};

inline void complex_op(bool b, bool & and, bool & or) {
	and = and && b;
	or = or || b;
}

// prior classification is done using mask values only:
// if all neighbours belong to one class then there is no other choice
classification prior(const mask & m, size_point const & c) {
	size_t w = m.width();
	size_t h = m.height();
	bool mul = true, sum = false;
	if (c.x > 0 && c.y > 0)
		complex_op(m(c.x-1, c.y-1), mul, sum);
	if (c.x > 0)
		complex_op(m(c.x-1, c.y  ), mul, sum);
	if (c.x > 0 && c.y + 1 < h)
		complex_op(m(c.x-1, c.y+1), mul, sum);
	if (c.y + 1 < h)
		complex_op(m(c.x  , c.y+1), mul, sum);
	if (c.x + 1 < w && c.y + 1 < h)
		complex_op(m(c.x+1, c.y+1), mul, sum);
	if (c.x + 1 < w)
		complex_op(m(c.x+1, c.y  ), mul, sum);
	if (c.x + 1 < w && c.y > 0)
		complex_op(m(c.x+1, c.y-1), mul, sum);
	if (c.y > 0)
		complex_op(m(c.x  , c.y-1), mul, sum);
	if (mul)
		return class_1;
	if (!sum)
		return class_0;
	return class_denial;
}

template <class R>
classification classify(sorted_neighbours<R> & sort_cache, const mask & m, 
	size_point const & c, icolor const & col, const params & p)
{
	assert (p.t > 1);
	uint nei_sorted_size = 2 * p.t;
	const neighbours & nei = sort_cache.get(c, col, nei_sorted_size);

	int d = 0;
	for (uint n = 0; n != nei.size(); ++n) {
		if (n+1 >= nei_sorted_size) {
			nei_sorted_size *= 2;
			sort_cache.get(c, col, nei_sorted_size);
		}
		const bool nei_state = m(c.x + nei[n].second.x, c.y + nei[n].second.y);
		if (nei_state)
			++d;
		else
			--d;

		//plateau (several equal color distances) process as single point
		if (n+1 != nei.size() && nei[n].first == nei[n+1].first)
			continue;

		if (d * d > p.t * (n+1))
			return (d > 0) ? class_1 : class_0;
	}
	return class_denial;
}

template <class R>
void core(
	const R & ref_im, const mask & ref_m, 
	const R & tgt_im, mask & tgt_m, 
	const affine * motion,
	const params & p, bool check_once
)
{
	size_extent ref_dim = ref_im.dim();
	assert(ref_dim == ref_m.dim());
	size_extent tgt_dim = tgt_im.dim();
	assert(tgt_dim == tgt_m.dim());
	if (ref_m.data() != tgt_m.data())
		check_once = true; //since classification of any pixel depends only from ref_m, which does not change
	sorted_neighbours<R> sort_cache05(ref_im,  5, p);
	sorted_neighbours<R> sort_cache10(ref_im, 10, p);
	sorted_neighbours<R> sort_cache20(ref_im, 20, p);
	sorted_neighbours<R> sort_cache40(ref_im, 40, p);

	if (!p.silent)
		clog << "refining initiated..." << endl;

	mask modifiable(tgt_dim);
	modifiable.fill(true);

	mask not_checked_yet;
	if (check_once) {
		not_checked_yet.resize(tgt_dim);
		not_checked_yet.fill(true);
	}

	mask next_m = tgt_m;

	affine back_proj;
	if (motion) {
		back_proj = motion->inverse();
		back_proj.translate(0.5f, 0.5f); //for proper rounding later
	} else
		back_proj.identity();
	int_rect ref_frame(0, 0, ref_dim.sx, ref_dim.sy);

	for (int pass = 0; ; ++pass) 
	{
		if (check_once)
			img::combine(modifiable, not_checked_yet, img::bit_and());
		if (p.modifiable)
			img::combine(modifiable, *p.modifiable, img::bit_and());

		uint checks = 0, changes = 0;//, scans = sort_cache.scans(), sorts = sort_cache.sorts();

		for (mask::iterator c = modifiable.begin(); c != modifiable.end(); ++c) {
			if (!*c)
				continue;
			bool tgt_b = tgt_m[c];
			if (p.change == turn_on_only && tgt_b || p.change == turn_off_only && !tgt_b) {
				*c = false;
				continue;
			}

			classification cl = prior(next_m, c);
			if (cl == class_denial) {
				++checks;
				if (check_once)
					not_checked_yet[c] = false;
				//back project c
				int_point ref_c = back_proj((size_point)c);
				if (ref_frame.inside(ref_c)) {
					cl = classify(sort_cache05, ref_m, ref_c, tgt_im[c], p);
					if (cl == class_denial)
						cl = classify(sort_cache10, ref_m, ref_c, tgt_im[c], p);
					if (cl == class_denial)
						cl = classify(sort_cache20, ref_m, ref_c, tgt_im[c], p);
					if (cl == class_denial)
						cl = classify(sort_cache40, ref_m, ref_c, tgt_im[c], p);
				}
				if (p.prefer != prefer_none && cl == class_denial) 
					cl = (p.prefer == prefer_masked_pixels) ? class_1 : class_0;
			}
			if (cl == class_0 && !tgt_b || cl == class_1 && tgt_b)
				cl = class_denial;
			if (cl == class_denial)
				*c = false;
			else {
				next_m[c] = !tgt_b;
				++changes;
			}
		}

/*		scans = sort_cache.scans() - scans;
		sorts = sort_cache.sorts() - sorts;
		if (!p.silent)
			clog << "pass " << pass <<
				": checks=" << checks << 
				", cache_size=" << sort_cache.size() <<
				", sorts=" << sorts << 
				", scans=" << scans << 
				", changes=" << changes << endl;*/

		if (changes == 0 //|| scans == 0 && sorts == 0 //sorts&scans are zero if some cycling takes place
			|| pass >= p.max_passes)
			break;

		tgt_m = next_m;

		// modified on the current pass points can influence state of other points
		// on the distance up to p.whs
		/*if (p.whs <= 2)
			img::dilate8(modifiable, p.whs); //for square window
		else
			img::dilate4(modifiable, p.whs); //for circle window (slightly less than needed)*/
		img::dilate4(modifiable, 40);
		sort_cache05.pass_done();
		sort_cache10.pass_done();
		sort_cache20.pass_done();
		sort_cache40.pass_done();
	}
}

////////////////////////
// Exported functions //
////////////////////////

void adjust(
	const byte_image & im, mask & m, 
	const params & p, bool check_once)
{
	core(im, m, im, m, 0, p, check_once);
}

void adjust(
	const argb_byte_image & im, mask & m, 
	const params & p, bool check_once)
{
	core(im, m, im, m, 0, p, check_once);
}

void track(
	const byte_image & ref_im, const mask & ref_m, 
	const byte_image & tgt_im, mask & tgt_m, 
	information_source isource,
	const affine * motion,
	const params & p
)
{
	if (isource == source_reference || isource == source_both) {
		if (!p.silent)
			clog << "\nrefine using reference frame\n";
		core(ref_im, ref_m, tgt_im, tgt_m, motion, p, true);
	}
	if (isource == source_target || isource == source_both) {
		if (!p.silent)
			clog << "\nrefine (adjust)\n";
		adjust(tgt_im, tgt_m, p, false);
	}
}

} //namespace refine3

} //namespace refine
